//! Connection pool implementation for HTTP proxy

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use log::debug;

/// Connection pool for managing reusable TCP connections
#[derive(Debug)]
pub struct ConnectionPool {
    connections: Arc<Mutex<HashMap<String, Vec<(TcpStream, Instant)>>>>,
    max_idle_time: Duration,
}

impl ConnectionPool {
    /// Create a new connection pool
    pub fn new() -> Self {
        Self {
            connections: Arc::new(Mutex::new(HashMap::new())),
            max_idle_time: Duration::from_secs(30),
        }
    }

    /// Create a new connection pool with custom idle timeout
    pub fn with_idle_timeout(timeout: Duration) -> Self {
        Self {
            connections: Arc::new(Mutex::new(HashMap::new())),
            max_idle_time: timeout,
        }
    }

    /// Get a connection from the pool or create a new one
    pub async fn get_or_create(&self, target_addr: &str) -> Result<TcpStream, std::io::Error> {
        // Try to get from pool first
        if let Some(stream) = self.get(target_addr).await {
            debug!("Reusing connection from pool for {}", target_addr);
            return Ok(stream);
        }

        // Create new connection
        debug!("Creating new connection to {}", target_addr);
        TcpStream::connect(target_addr).await
    }

    /// Get a connection from the pool if available
    pub async fn get(&self, target_addr: &str) -> Option<TcpStream> {
        let mut pool = self.connections.lock().await;
        
        if let Some(connections) = pool.get_mut(target_addr) {
            // Clean up expired connections
            let now = Instant::now();
            connections.retain(|(_, instant)| now.duration_since(*instant) < self.max_idle_time);
            
            if let Some((stream, _)) = connections.pop() {
                return Some(stream);
            }
        }
        
        None
    }

    /// Return a connection to the pool
    pub async fn put(&self, target_addr: String, stream: TcpStream) {
        let mut pool = self.connections.lock().await;
        let connections = pool.entry(target_addr).or_default();
        connections.push((stream, Instant::now()));
        debug!("Returned connection to pool, total: {}", connections.len());
    }

    /// Clear all connections from the pool
    pub async fn clear(&self) {
        let mut pool = self.connections.lock().await;
        pool.clear();
    }

    /// Get the number of connections in the pool
    pub async fn len(&self) -> usize {
        let pool = self.connections.lock().await;
        pool.values().map(|v| v.len()).sum()
    }

    /// Check if the pool is empty
    pub async fn is_empty(&self) -> bool {
        self.len().await == 0
    }
}

impl Default for ConnectionPool {
    fn default() -> Self {
        Self::new()
    }
}